{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "from PIL import Image, ImageDraw\n",
    "import tensorflow as tf\n",
    "import json\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "import sys\n",
    "sys.path.append('/home/dan/work/cocoapi/PythonAPI/')\n",
    "\n",
    "from pycocotools.coco import COCO\n",
    "from pycocotools.cocoeval import COCOeval\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from detector import Detector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_PATH = 'model.pb'\n",
    "detector = Detector(MODEL_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load COCO groundtruth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGES_DIR = '/home/dan/datasets/COCO/images/'\n",
    "ANNOTATIONS_DIR = '/home/dan/datasets/COCO/annotations/'\n",
    "DATA_TYPE = 'val2017'  # train2017 or val2017"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coco = COCO(os.path.join(ANNOTATIONS_DIR, 'instances_{}.json'.format(DATA_TYPE)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../data/coco_labels.txt', 'r') as f:\n",
    "    labels = {line.strip(): i for i, line in enumerate(f.readlines()) if line.strip()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "categories = coco.loadCats(coco.getCatIds())\n",
    "class_to_coco_id = {c['name']: c['id']  for c in categories}\n",
    "\n",
    "# my integer encoding (from 0 to 79) to official integer encoding\n",
    "integer_to_coco_id = {labels[n]: i for n, i in class_to_coco_id.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compute predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "catIds = coco.getCatIds()\n",
    "imgIds = coco.getImgIds()\n",
    "len(imgIds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "results = []\n",
    "\n",
    "for image_id in tqdm(imgIds):\n",
    "\n",
    "    image_metadata = coco.loadImgs(image_id)[0]\n",
    "    image_path = os.path.join(IMAGES_DIR, DATA_TYPE, image_metadata['file_name'])\n",
    "    image = cv2.imread(image_path)\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    height, width, _ = image.shape\n",
    "\n",
    "    boxes, labels, scores = detector(image, score_threshold=0.15)\n",
    "    scaler = np.array([height, width, height, width], dtype='float32')\n",
    "    boxes = boxes * scaler\n",
    "\n",
    "    for i in range(len(boxes)):\n",
    "        ymin, xmin, ymax, xmax = boxes[i]\n",
    "\n",
    "        x, y = int(xmin), int(ymin)\n",
    "        w, h = int(xmax - xmin), int(ymax - ymin)\n",
    "\n",
    "        results.append({\n",
    "            \"image_id\" : int(image_metadata['id']), \"category_id\" : int(integer_to_coco_id[labels[i]]), \n",
    "            \"bbox\" : [x, y, w, h], \"score\" : float(scores[i])\n",
    "        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open('coco_predictions.json', 'w') as f:\n",
    "    json.dump(results, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show detections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coco_detections = coco.loadRes('coco_predictions.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_random():\n",
    "    \n",
    "    i = np.random.randint(0, len(imgIds))\n",
    "    image_id = imgIds[i]\n",
    "    image_metadata = coco.loadImgs(image_id)[0]\n",
    "    \n",
    "    image_path = os.path.join(IMAGES_DIR, DATA_TYPE, image_metadata['file_name'])\n",
    "    image = cv2.imread(image_path)\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    annIds = coco_detections.getAnnIds(imgIds=image_metadata['id'], catIds=catIds)\n",
    "    annotations = coco_detections.loadAnns(annIds)\n",
    "    return image, annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image, annotations = get_random()\n",
    "\n",
    "plt.figure(dpi=100, figsize=(17, 10))\n",
    "plt.imshow(image)\n",
    "plt.axis('off')\n",
    "coco.showAnns(annotations);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cocoEval = COCOeval(cocoGt=coco, cocoDt=coco_detections, iouType='bbox')\n",
    "cocoEval.params.imgIds = imgIds\n",
    "cocoEval.params.catIds = catIds\n",
    "cocoEval.evaluate()\n",
    "cocoEval.accumulate()\n",
    "cocoEval.summarize()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
